from typing import List, Optional, Any
from .lexer import Token, Lexer
from .ast_nodes import *

class Parser:
    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.pos = 0

    def parse(self) -> Program:
        definitions = []
        while self.current().type != 'EOF':
            definitions.append(self.parse_definition())
        return Program(definitions)

    def parse_definition(self) -> ASTNode:
        if self.match_keyword('设'):
            return self.parse_function_or_variable_def()
        elif self.match_keyword('记录'):
            return self.parse_record_def()
        else:
            # For now, only support top-level definitions starting with '设' or '记录'
            raise SyntaxError(f"Unexpected token at top level: {self.current()}")

    def parse_record_def(self) -> RecordDef:
        # 记录 Name
        name_token = self.consume('IDENTIFIER')
        name = name_token.value
        
        # 期待 INDENT 开始字段列表
        if not self.match('INDENT'):
            raise SyntaxError(f"记录定义 '{name}' 需要代码块")
        
        fields = []
        while not self.check('DEDENT') and not self.check('EOF'):
            # 解析字段: 字段名 : 类型
            field_name = self.consume('IDENTIFIER').value
            self.consume('COLON')
            field_type = self.parse_type()
            fields.append(Param(field_name, field_type))
        
        self.consume('DEDENT')
        return RecordDef(name, fields)

    def parse_function_or_variable_def(self) -> ASTNode:
        # 设 name ...
        name_token = self.consume('IDENTIFIER')
        name = name_token.value

        # 检查是函数还是变量
        # 如果下一个 token 是 '=', 它是变量定义： 设 x = 1
        # 如果下一个 token 是 ':', 它可能是变量有类型： 设 x : Int = 1
        # 如果下一个 token 是 identifier 或 '(', 它是函数： 设 f x = ... 或 设 f (x: Int) ...
        
        if self.match('ASSIGN'):
            # 变量定义，没有类型
            value = self.parse_expression()
            return VariableDef(name, value)
        
        elif self.match('COLON'):
            # 变量定义有类型？或函数有返回类型（不太可能 设 name : Type）
            # 假设 设 name : Type = Expr 为变量
            type_node = self.parse_type()
            if self.match('ASSIGN'):
                value = self.parse_expression()
                return VariableDef(name, value, type_node)
            else:
                 # 可能是函数定义签名： 设 main : Int
                 # 对于现在，让我们假设它是一个函数，如果立即没有 '='
                 pass

        # 函数定义
        params = []
        while not self.check('ASSIGN') and not self.check('COLON') and not self.check('INDENT') and not self.check('EOF'):
            params.append(self.parse_param())
        
        return_type = None
        if self.match('COLON'):
            return_type = self.parse_type()
        
        # 函数体
        # 可以是内联： 设 f x = x + 1
        # 或块：
        # 设 f x
        #   x + 1
        
        if self.match('ASSIGN'):
            # 内联体，包装在 Block
            expr = self.parse_expression()
            body = Block([expr])
        elif self.match('INDENT'):
            body = self.parse_block()
        else:
            raise SyntaxError(f"函数 “{name}” 需要 “=” 或代码块")
            
        return FunctionDef(name, params, return_type, body)

    def parse_param(self) -> Param:
        if self.match('LPAREN'):
            # (name : type)
            name = self.consume('IDENTIFIER').value
            self.consume('COLON')
            type_node = self.parse_type()
            self.consume('RPAREN')
            return Param(name, type_node)
        else:
            # name
            name = self.consume('IDENTIFIER').value
            return Param(name, None)

    def parse_type(self) -> TypeNode:
        # 简单类型解析
        name = self.consume('IDENTIFIER', 'KEYWORD').value # 类型可以是关键词，如 '数'
        # 处理泛型： List[Int] -> 列[数]
        generics = []
        if self.match('LBRACKET'):
            while not self.check('RBRACKET'):
                generics.append(self.parse_type())
                if not self.check('RBRACKET'):
                    self.consume('COMMA')
            self.consume('RBRACKET')
        return TypeNode(name, generics)

    def parse_block(self) -> Block:
        statements = []
        while not self.check('DEDENT') and not self.check('EOF'):
            statements.append(self.parse_statement())
        self.consume('DEDENT')
        return Block(statements)

    def parse_statement(self) -> ASTNode:
        # 语句可以是定义（局部变量）或表达式
        if self.match_keyword('设'):
            return self.parse_function_or_variable_def()
        elif self.match_keyword('返回'):
            # 显式返回
            expr = self.parse_expression()
            # 如果需要，包装在 ReturnNode 中，或现在就当 Expr
            # InkLang 是基于表达式的，'返回' 可能只是用于早期退出或清晰
            return expr 
        else:
            return self.parse_expression()

    def parse_expression(self) -> Expr:
        return self.parse_binary_op(0)

    def parse_binary_op(self, min_prec: int) -> Expr:
        # 优先级表
        PRECEDENCE = {
            'ASSIGN': 1,
            'PIPE': 2,
            'OR': 3,
            'AND': 4,
            'EQ': 5, 'NE': 5, 'LT': 5, 'GT': 5, 'LE': 5, 'GE': 5,
            'PLUS': 6, 'MINUS': 6,
            'MUL': 7, 'DIV': 7, 'MOD': 7,
            'DOT': 8
        }
        
        left = self.parse_atom()
        
        while True:
            token = self.current()
            if token.type not in PRECEDENCE or PRECEDENCE[token.type] < min_prec:
                break
            
            op = token.type
            self.advance()
            if op == 'DOT':
                # Field Access: left.right
                # right must be an Identifier
                token = self.consume('IDENTIFIER')
                left = FieldAccess(left, token.value)
            else:
                right = self.parse_binary_op(PRECEDENCE[op] + 1)
                left = BinOp(left, token.value, right)
            
        return left

    def parse_atom(self) -> Expr:
        if self.match_keyword('若'):
            return self.parse_if_expr()
        
        token = self.current()
        
        if token.type == 'NUMBER':
            self.advance()
            return Literal(token.value, '数') # 简化，可能是 float
        elif token.type == 'STRING':
            self.advance()
            return Literal(token.value, '词')
        elif token.type == 'LBRACKET':
            # 列表字面量： [1, 2, 3]
            return self.parse_list_literal()
        elif token.type == 'IDENTIFIER':
            self.advance()
            expr = Identifier(token.value)
            
            # Check for Record Literal: Type { ... }
            if self.match('LBRACE'):
                # Record Literal
                field_values = {}
                while not self.check('RBRACE') and not self.check('EOF'):
                    field_name = self.consume('IDENTIFIER').value
                    self.consume('ASSIGN')
                    value = self.parse_expression()
                    field_values[field_name] = value
                    if not self.check('RBRACE'):
                        self.consume('COMMA')
                self.consume('RBRACE')
                return RecordLiteral(token.value, field_values)
            
            # Check for function call: f a b
            # If next is atom start, it's a call
            while self.is_atom_start():
                # 检查是否在同一行
                # 如果换行了，不再视为函数调用参数（除非显式续行，暂不支持）
                current_token = self.current()
                prev_token = self.tokens[self.pos - 1]
                if current_token.line > prev_token.line:
                    break
                
                arg = self.parse_atom()
                expr = CallExpr(expr, [arg])
            return expr
        elif token.type == 'LPAREN':
            self.advance()
            # 可以是元组 (a, b, c) 或分组表达式 (expr)
            if self.check('RPAREN'):
                # 空元组 ()
                self.advance()
                return TupleLiteral([])
            
            first_expr = self.parse_expression()
            if self.match('COMMA'):
                # 元组
                elements = [first_expr]
                elements.append(self.parse_expression())
                while self.match('COMMA'):
                    elements.append(self.parse_expression())
                self.consume('RPAREN')
                return TupleLiteral(elements)
            else:
                # 分组表达式
                self.consume('RPAREN')
                return first_expr
        else:
            # 检查布尔字面量
            if token.type == 'KEYWORD' and token.value in ('真', '假'):
                self.advance()
                return Literal(token.value, '两仪')
            
            raise SyntaxError(f"意外的 token 在表达式中: {token}")

    def parse_list_literal(self) -> ListLiteral:
        # 已消费 '['
        self.consume('LBRACKET')
        elements = []
        while not self.check('RBRACKET') and not self.check('EOF'):
            elements.append(self.parse_expression())
            if not self.check('RBRACKET'):
                self.consume('COMMA')
        self.consume('RBRACKET')
        return ListLiteral(elements)

    def parse_if_expr(self) -> IfExpr:
        # 若 cond 则 then [否则 else]
        condition = self.parse_expression()
        self.consume_keyword('则')
        
        if self.match('INDENT'):
            then_branch = self.parse_block()
        else:
            # 内联 then
            expr = self.parse_expression()
            then_branch = Block([expr])
            
        else_branch = None
        if self.match_keyword('否则'):
            if self.match('INDENT'):
                else_branch = self.parse_block()
            elif self.check_keyword('若'):
                else_branch = self.parse_if_expr() # else if
            else:
                expr = self.parse_expression()
                else_branch = Block([expr])
                
        return IfExpr(condition, then_branch, else_branch)

    def is_atom_start(self) -> bool:
        t = self.current()
        return t.type in ('NUMBER', 'STRING', 'IDENTIFIER', 'LPAREN', 'LBRACKET') or \
               (t.type == 'KEYWORD' and t.value in ('真', '假', '若'))

    # 辅助方法
    def current(self) -> Token:
        if self.pos >= len(self.tokens):
            return Token('EOF', '', -1, -1)
        return self.tokens[self.pos]

    def advance(self):
        self.pos += 1

    def check(self, type: str) -> bool:
        return self.current().type == type

    def check_keyword(self, value: str) -> bool:
        return self.current().type == 'KEYWORD' and self.current().value == value

    def match(self, type: str) -> bool:
        if self.check(type):
            self.advance()
            return True
        return False

    def match_keyword(self, value: str) -> bool:
        if self.check_keyword(value):
            self.advance()
            return True
        return False

    def consume(self, *types: str) -> Token:
        if self.current().type in types:
            t = self.current()
            self.advance()
            return t
        raise SyntaxError(f"期望 {types}，但得到 {self.current()}")

    def consume_keyword(self, value: str) -> Token:
        if self.check_keyword(value):
            t = self.current()
            self.advance()
            return t
        raise SyntaxError(f"期望关键词 '{value}'，但得到 {self.current()}")
